#include <x86.h>

#define PG_ALIGN __attribute((aligned(PGSIZE)))

static PDE kpdirs[NR_PDE] PG_ALIGN;
static PTE kptabs[PMEM_SIZE / PGSIZE] PG_ALIGN;
static void* (*palloc_f)();
static void (*pfree_f)(void*);

_Area segments[] = {      // Kernel memory mappings
  {.start = (void*)0,          .end = (void*)PMEM_SIZE}
};

#define NR_KSEG_MAP (sizeof(segments) / sizeof(segments[0]))

void _pte_init(void* (*palloc)(), void (*pfree)(void*)) {
  palloc_f = palloc;
  pfree_f = pfree;

  int i;

  // make all PDEs invalid
  for (i = 0; i < NR_PDE; i ++) {
    kpdirs[i] = 0;
  }

  PTE *ptab = kptabs;
  for (i = 0; i < NR_KSEG_MAP; i ++) {
    uint32_t pdir_idx = (uintptr_t)segments[i].start / (PGSIZE * NR_PTE);
    uint32_t pdir_idx_end = (uintptr_t)segments[i].end / (PGSIZE * NR_PTE);
    for (; pdir_idx < pdir_idx_end; pdir_idx ++) {
      // fill PDE
      kpdirs[pdir_idx] = (uintptr_t)ptab | PTE_P;

      // fill PTE
      PTE pte = PGADDR(pdir_idx, 0, 0) | PTE_P;
      PTE pte_end = PGADDR(pdir_idx + 1, 0, 0) | PTE_P;
      for (; pte < pte_end; pte += PGSIZE) {
        *ptab = pte;
        ptab ++;
      }
    }
  }

  set_cr3(kpdirs);
  set_cr0(get_cr0() | CR0_PG);
}

void _protect(_Protect *p) {
  PDE *updir = (PDE*)(palloc_f());
  p->ptr = updir;
  // map kernel space
  for (int i = 0; i < NR_PDE; i ++) {
    updir[i] = kpdirs[i];
  }

  p->area.start = (void*)0x8000000;
  p->area.end = (void*)0xc0000000;
}

void _release(_Protect *p) {
}

void _switch(_Protect *p) {
  set_cr3(p->ptr);
}

void _map(_Protect *p, void *va, void *pa) {
  // 1. 获取页目录基地址（CR3寄存器值，物理地址）
  PDE *dir_base = (PDE *)(p->ptr);

  // 2. 计算虚拟地址的页目录索引（高10位）和页表索引（中间10位）
  uint32_t page = PTX(va);  // (va >> 12) & 0x3FF
  uint32_t dir = PDX(va);   // (va >> 22) & 0x3FF

  // 3. 检查页目录项是否存在（Present位是否为1）
  if (!(dir_base[dir] & 0x1)) {
      // 4. 若不存在，申请一页物理内存作为新页表
      PTE *uptab = (PTE *)(palloc_f());
      // 5. 设置页目录项：物理地址（高20位） + Present位（0x1）
      dir_base[dir]  = (uint32_t)uptab | PTE_P;
  }

  // 6. 获取页表基地址（清除页目录项的低12位标志）
  PTE *page_base = (PTE *)(dir_base[dir] & 0xfffff000);

  // 7. 定位到页表项（根据页表索引）
  PTE *pte = &page_base[page];

  // 8. 设置页表项：物理地址（高20位） + Present位（0x1）
  *pte = (uint32_t)pa | PTE_P;
}

void _unmap(_Protect *p, void *va) {
}

_RegSet *_umake(_Protect *p, _Area ustack, _Area kstack, void *entry, char *const argv[], char *const envp[]) {
  uint32_t *ptr = ustack.end;

  // --------------------- 参数区 ---------------------
  *ptr-- = 0; // envp
  *ptr-- = 0; // argv
  *ptr-- = 0; // argc

  // --------------------- 陷阱帧 ---------------------
  // 硬件自动压入部分
  *ptr-- = 0x2 | (1<<9);  // eflags (IF=1)
  *ptr-- = 0x8;           // cs
  *ptr-- = (uint32_t)entry; // eip
  *ptr-- = 0x0;           // error code
  *ptr-- = 0x81;          // irq id (trap)
  
  // pushal保存的通用寄存器
  for(int i = 0; i < 8; i++) 
    *ptr-- = 0; // edi,esi,ebp,esp,ebx,edx,ecx,eax

  return (_RegSet*)(ptr+1); // 陷阱帧起始地址
}
